/* Filename: mytbf.c
 * Description: 查询法/通知法实现令牌桶的库
 * Last modified: humble 2020-09-03 21:12:35
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <pthread.h>
#include <sched.h>
#include "mytbf.h"

#define mutex_only (0) //查询法
#define with_cond (1) //通知法

#define demo (with_cond)

#define MYTBF_MAX (1024) //该库最多支持桶的数量

//单个桶的属性
typedef struct{
    int cps; //char per second 速率
    int burst; //桶最大容量值
    int token; //当前桶内有的字符总数
    int pos; //桶在库数组的下标
    pthread_mutex_t mut; //自己桶内token的互斥锁
    pthread_cond_t cond; //自己同内token的条件锁
}mytbf_st;

static mytbf_st *job[MYTBF_MAX]; //库数组
static pthread_mutex_t mut_job = PTHREAD_MUTEX_INITIALIZER; //库数组锁

//时钟线程，定时给库数组的每个桶加cps个字符
//当库数组里有令牌桶时该线程被创建
//当库数组里没令牌桶时该线程被回收
static pthread_t tid_alrm;

static pthread_once_t init_once = PTHREAD_ONCE_INIT; //控制执行一次的变量

static void *thr_alrm(void *p)
{
    int i;
    while(1){
        pthread_mutex_lock(&mut_job);//每秒锁一次库数组，遍历给所有桶加它对应的cps个字符
        for(i = 0; i < MYTBF_MAX; i++)
        {
            if(job[i]){
                pthread_mutex_lock(&job[i]->mut);
                job[i]->token += job[i]->cps;
                if(job[i]->token > job[i]->burst){ job[i]->token = job[i]->burst; }//防止超过桶容量上限
#if (demo == with_cond)
                pthread_cond_broadcast(&job[i]->cond); //添加完token后，唤醒全部线程过来抢
#endif
                pthread_mutex_unlock(&job[i]->mut);
            }
        }
        pthread_mutex_unlock(&mut_job);
        sleep(1);
    }
    pthread_exit(NULL);
}

//unload模块没加锁，只能由1个人调用
//如果2个人调用就有一个人重复释放资源（即使加锁也无法解决）
static void module_unload(void)
{
    int i;
    pthread_cancel(tid_alrm); //给时钟线程发送一个线程取消请求
    pthread_join(tid_alrm, NULL); //为时钟线程收尸
    for(i = 0; i < MYTBF_MAX; i++)
    {
        if(job[i]){ mytbf_destroy(job[i]); }
    }
    pthread_mutex_destroy(&mut_job);
}

static void module_load(void)
{
    int err;
    err = pthread_create(&tid_alrm, NULL, thr_alrm, NULL);
    if(err){ fprintf(stderr, "pthread_create():%s\n", strerror(err)); exit(1); }
    atexit(module_unload);
}

/*
* 为了不破坏调用者对令牌桶操作的原子性，
* 在该函数内加锁可能会导致死锁，
* 所以该函数内部无法加锁，
* 必须在调用该函数之前先加锁。
*/
static int get_free_pos_unlocked(void)
{
    int i;
    for(i = 0; i < MYTBF_MAX; i++)
    {
        if(!job[i]){ return i; }
    }
    return -1;
}

mytbf_t *mytbf_init(int cps, int burst)
{
    mytbf_st *me;
    int pos;
    if(cps <= 0 || burst <= 0){ return NULL; }

#if 0
    lock() //第三把锁(建议) 或 库数组锁(但是会暂时影响别人访问库数组)
    if(!inited){ inited = 1; module_load(); }
    unlock()
#else //上下代码等价
    pthread_once(&init_once, module_load);//只有初次调用mytbf_init，module_load才会被调用
#endif

    //申请一个桶的资源并初始化
    if(!(me = malloc(sizeof(mytbf_st)))){ return NULL; }
    me->token = 0;
    me->cps = cps;
    me->burst = burst;
    pthread_mutex_init(&me->mut, NULL);
#if (demo == with_cond)
    pthread_cond_init(&me->cond, NULL);
#endif

    //抢锁，把桶放入数组，解锁
    pthread_mutex_lock(&mut_job);
    pos = get_free_pos_unlocked();
    if(pos < 0 || pos >= MYTBF_MAX){
        pthread_mutex_unlock(&mut_job);
        free(me);
        return NULL;
    }
    me->pos = pos;
    job[pos] = me;
    pthread_mutex_unlock(&mut_job);

    return me;
}

//从ptr这个桶内取size
int mytbf_fetch(mytbf_t *ptr, int size)
{
    mytbf_st *me = (mytbf_st *)ptr;
    if((!me) || (size <= 0)){ return -EINVAL; }

    pthread_mutex_lock(&me->mut);
    while(me->token <= 0){ //当前桶内没有字符
#if (demo == mutex_only)
        pthread_mutex_unlock(&me->mut);
        sched_yield();
        pthread_mutex_lock(&me->mut);
#elif (demo == with_cond) //上下代码等价(但下面这句会休眠)
        pthread_cond_wait(&me->cond, &me->mut); //被唤醒后，开始抢锁mut(抢到并发现桶内有字符就跳出while，否则就回到这里休眠)
#endif
    }
    //此时桶内终于有字符了
    //取字符
    int n = me->token < size ? me->token : size;
    me->token -= n;
    pthread_mutex_unlock(&me->mut);

    return n;
}

//程序用剩size个字符，归还给ptr这个桶
int mytbf_return(mytbf_t *ptr, int size)
{
    mytbf_st *me = (mytbf_st *)ptr;
    if((!me) || (size <= 0)){ return -EINVAL; }

    pthread_mutex_lock(&me->mut);
    me->token += size;
    if(me->token > me->burst){ me->token = me->burst; } //防止超过桶容量上限
#if (demo == with_cond)
    pthread_cond_broadcast(&me->cond);
#endif
    pthread_mutex_unlock(&me->mut);

    return size;
}

//销毁一个桶
int mytbf_destroy(mytbf_t *ptr)
{
    mytbf_st *me = (mytbf_st *)ptr;
    if(!me){ return -1; }

    pthread_mutex_lock(&mut_job);
    job[me->pos] = NULL;
    pthread_mutex_unlock(&mut_job);

    pthread_mutex_destroy(&me->mut);
#if (demo == with_cond)
    pthread_cond_destroy(&me->cond);
#endif
    free(me);
    return 0;
}
